use motore::{layer::Layer, service::Service};
use volo::context::Context;

use crate::{
    context::client::Config,
    error::ClientError,
    request::{Request, RequestPartsExt},
};

/// [`Layer`] for applying timeout from [`Config`].
///
/// This layer will be applied by default when using [`ClientBuilder::build`], without this layer,
/// timeout from [`Client`] or [`CallOpt`] will not work.
///
/// [`Client`]: crate::client::Client
/// [`ClientBuilder::build`]: crate::client::ClientBuilder::build
/// [`CallOpt`]: crate::client::CallOpt
#[derive(Clone, Debug, Default)]
pub struct Timeout;

impl<S> Layer<S> for Timeout {
    type Service = TimeoutService<S>;

    fn layer(self, inner: S) -> Self::Service {
        TimeoutService { inner }
    }
}

/// The [`Service`] generated by [`Timeout`].
///
/// See [`Timeout`] for more details.
pub struct TimeoutService<S> {
    inner: S,
}

impl<Cx, B, S> Service<Cx, Request<B>> for TimeoutService<S>
where
    Cx: Context<Config = Config> + Send,
    B: Send,
    S: Service<Cx, Request<B>, Error = ClientError> + Send + Sync,
{
    type Response = S::Response;
    type Error = S::Error;

    async fn call(&self, cx: &mut Cx, req: Request<B>) -> Result<Self::Response, Self::Error> {
        let timeout = cx.rpc_info().config().timeout().cloned();

        if let Some(duration) = timeout {
            let url = req.url();
            let fut = self.inner.call(cx, req);
            let sleep = tokio::time::sleep(duration);

            tokio::select! {
                res = fut => res,
                _ = sleep => {
                    if let Some(url) = url {
                        tracing::warn!("[Volo-HTTP] request timeout on `{url}`");
                    }
                    Err(crate::error::client::timeout().with_endpoint(cx.rpc_info().callee()))
                }
            }
        } else {
            self.inner.call(cx, req).await
        }
    }
}
